package com.vwa.DeepSeekR1.Controller;

import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.apache.commons.lang.StringUtils;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

@RestController
public class ChatCompletionsController {

    static String apiKey = "a432402f-d1d2-40a5-9e24-7ac83687d2d5";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    // The output time of the reasoning model is relatively long. Please increase the timeout period.
    static ArkService service = ArkService.builder().timeout(Duration.ofSeconds(1800)).connectTimeout(Duration.ofSeconds(20)).dispatcher(dispatcher).connectionPool(connectionPool).baseUrl("https://ark.cn-beijing.volces.com/api/v3").apiKey(apiKey).build();

    @GetMapping("/getChatResponse")
    public ResponseEntity<String> getChatResponse(@RequestParam("content") String content) {
        System.out.println("Received content: " + content);

        // 创建消息列表
        final List<ChatMessage> streamMessages = new ArrayList<>();
        final ChatMessage streamUserMessage = ChatMessage.builder()
                .role(ChatMessageRole.USER)
                .content(content)
                .build();
        streamMessages.add(streamUserMessage);

        // 创建请求对象
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model("deepseek-r1-250120")
                .messages(streamMessages)
                .build();

        // 用于拼接流式响应内容
        StringBuilder responseBuilder = new StringBuilder();

        // 使用 CompletableFuture 等待流式响应完成
        CompletableFuture<String> responseFuture = new CompletableFuture<>();

        service.streamChatCompletion(streamChatCompletionRequest)
                .doOnError(error -> {
                    error.printStackTrace();  // 打印错误信息
                    responseFuture.completeExceptionally(error);  // 完成 Future 并传递异常
                })
                .subscribe(delta -> {
                    if (delta != null && !delta.getChoices().isEmpty()) {
                        //String reasoning = delta.getChoices().get(0).getMessage().getReasoningContent();
                        //System.out.print(reasoning);
                        String reasoningContent = delta.getChoices().get(0).getMessage().getContent().toString();
                        if (StringUtils.isNotEmpty(reasoningContent)) {
                            responseBuilder.append(reasoningContent);  // 拼接响应内容
                            System.out.print(reasoningContent);  // 打印到控制台
                        }
                    }
                }, responseFuture::completeExceptionally, () -> {
                    responseFuture.complete(responseBuilder.toString());  // 完成 Future 并返回结果
                });

        try {
            // 等待最多50秒获取响应
            String finalResponse = responseFuture.get(50, TimeUnit.SECONDS);
            return ResponseEntity.ok(finalResponse);
        } catch (Exception e) {
            e.printStackTrace();  // 打印异常堆栈
            return ResponseEntity.status(500)
                    .body("Error while processing the chat request: " + e.getMessage());
        }
    }

}